#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                            QListWidget, QListWidgetItem, QLabel, QComboBox,
                            QToolBar, QAction, QFileDialog, QSplitter, 
                            QScrollArea, QFrame, QGroupBox, QFormLayout, 
                            QLineEdit, QSpinBox, QDoubleSpinBox, QCheckBox,
                            QTabWidget, QTextEdit, QSlider, QRadioButton,
                            QButtonGroup, QMessageBox, QProgressBar, QGridLayout)
from PyQt5.QtCore import Qt, QSize, QRect, QPoint, pyqtSignal
from PyQt5.QtGui import QIcon, QPixmap, QFont

class ConfigPanel(QWidget):
    """模型配置面板"""
    
    # 自定义信号
    config_saved = pyqtSignal(str)  # 配置保存信号
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        
    def init_ui(self):
        """初始化UI"""
        # 创建主布局
        main_layout = QVBoxLayout(self)
        
        # 创建工具栏
        toolbar = QToolBar()
        
        # 添加工具栏按钮
        new_config_action = QAction("新建配置", self)
        new_config_action.triggered.connect(self.new_config)
        toolbar.addAction(new_config_action)
        
        load_config_action = QAction("加载配置", self)
        load_config_action.triggered.connect(self.load_config)
        toolbar.addAction(load_config_action)
        
        save_config_action = QAction("保存配置", self)
        save_config_action.triggered.connect(self.save_config)
        toolbar.addAction(save_config_action)
        
        main_layout.addWidget(toolbar)
        
        # 创建水平分割器
        splitter = QSplitter(Qt.Horizontal)
        
        # 左侧配置列表
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        
        # 添加配置列表组
        configs_group = QGroupBox("已保存配置")
        configs_layout = QVBoxLayout(configs_group)
        
        self.config_list = QListWidget()
        self.config_list.itemClicked.connect(self.select_config)
        configs_layout.addWidget(self.config_list)
        
        # 添加一些示例配置
        self.config_list.addItem("YOLOv10n - 默认配置")
        self.config_list.addItem("YOLOv10s - 默认配置")
        self.config_list.addItem("YOLOv10m - 默认配置")
        self.config_list.addItem("自定义配置 1")
        
        left_layout.addWidget(configs_group)
        
        # 添加配置预览
        preview_group = QGroupBox("配置预览")
        preview_layout = QVBoxLayout(preview_group)
        
        self.preview_text = QTextEdit()
        self.preview_text.setReadOnly(True)
        preview_layout.addWidget(self.preview_text)
        
        left_layout.addWidget(preview_group)
        
        # 创建右侧配置面板
        right_panel = QScrollArea()
        right_panel.setWidgetResizable(True)
        
        config_widget = QWidget()
        config_layout = QVBoxLayout(config_widget)
        
        # 使用标签页组织配置选项
        tabs = QTabWidget()
        
        # 基本配置标签页
        basic_tab = QWidget()
        basic_layout = QFormLayout(basic_tab)
        
        # 模型选择
        self.model_type_combo = QComboBox()
        self.model_type_combo.addItems(["YOLOv10n", "YOLOv10s", "YOLOv10m", "YOLOv10l", "YOLOv10x"])
        self.model_type_combo.currentIndexChanged.connect(self.update_preview)
        basic_layout.addRow("模型类型:", self.model_type_combo)
        
        # 输入尺寸
        self.input_size_combo = QComboBox()
        self.input_size_combo.addItems(["640x640", "512x512", "384x384", "1280x1280"])
        self.input_size_combo.setEditable(True)
        self.input_size_combo.currentIndexChanged.connect(self.update_preview)
        basic_layout.addRow("输入尺寸:", self.input_size_combo)
        
        # 批次大小
        self.batch_size_spin = QSpinBox()
        self.batch_size_spin.setRange(1, 128)
        self.batch_size_spin.setValue(16)
        self.batch_size_spin.valueChanged.connect(self.update_preview)
        basic_layout.addRow("批次大小:", self.batch_size_spin)
        
        # 训练轮数
        self.epochs_spin = QSpinBox()
        self.epochs_spin.setRange(1, 1000)
        self.epochs_spin.setValue(100)
        self.epochs_spin.valueChanged.connect(self.update_preview)
        basic_layout.addRow("训练轮数:", self.epochs_spin)
        
        # 设备选择
        self.device_combo = QComboBox()
        self.device_combo.addItems(["cuda:0", "cpu"])
        self.device_combo.currentIndexChanged.connect(self.update_preview)
        basic_layout.addRow("训练设备:", self.device_combo)
        
        tabs.addTab(basic_tab, "基本设置")
        
        # 优化器配置标签页
        optimizer_tab = QWidget()
        optimizer_layout = QFormLayout(optimizer_tab)
        
        # 优化器选择
        self.optimizer_combo = QComboBox()
        self.optimizer_combo.addItems(["SGD", "Adam", "AdamW"])
        self.optimizer_combo.currentIndexChanged.connect(self.update_preview)
        optimizer_layout.addRow("优化器:", self.optimizer_combo)
        
        # 学习率
        self.lr_spin = QDoubleSpinBox()
        self.lr_spin.setRange(0.00001, 0.1)
        self.lr_spin.setValue(0.01)
        self.lr_spin.setDecimals(5)
        self.lr_spin.setSingleStep(0.001)
        self.lr_spin.valueChanged.connect(self.update_preview)
        optimizer_layout.addRow("学习率:", self.lr_spin)
        
        # 学习率策略
        self.lr_scheduler_combo = QComboBox()
        self.lr_scheduler_combo.addItems(["cosine", "linear", "step"])
        self.lr_scheduler_combo.currentIndexChanged.connect(self.update_preview)
        optimizer_layout.addRow("学习率策略:", self.lr_scheduler_combo)
        
        # 权重衰减
        self.weight_decay_spin = QDoubleSpinBox()
        self.weight_decay_spin.setRange(0.0, 0.1)
        self.weight_decay_spin.setValue(0.0005)
        self.weight_decay_spin.setDecimals(5)
        self.weight_decay_spin.setSingleStep(0.0001)
        self.weight_decay_spin.valueChanged.connect(self.update_preview)
        optimizer_layout.addRow("权重衰减:", self.weight_decay_spin)
        
        tabs.addTab(optimizer_tab, "优化器设置")
        
        # 数据增强配置标签页
        augmentation_tab = QWidget()
        augmentation_layout = QVBoxLayout(augmentation_tab)
        
        # 常用增强设置
        common_aug_group = QGroupBox("常用增强")
        common_aug_layout = QFormLayout(common_aug_group)
        
        # 水平翻转
        self.flip_h_check = QCheckBox()
        self.flip_h_check.setChecked(True)
        self.flip_h_check.stateChanged.connect(self.update_preview)
        common_aug_layout.addRow("水平翻转:", self.flip_h_check)
        
        # 垂直翻转
        self.flip_v_check = QCheckBox()
        self.flip_v_check.setChecked(False)
        self.flip_v_check.stateChanged.connect(self.update_preview)
        common_aug_layout.addRow("垂直翻转:", self.flip_v_check)
        
        # 随机旋转
        self.rotate_check = QCheckBox()
        self.rotate_check.setChecked(True)
        self.rotate_check.stateChanged.connect(self.update_preview)
        common_aug_layout.addRow("随机旋转:", self.rotate_check)
        
        # 马赛克增强
        self.mosaic_check = QCheckBox()
        self.mosaic_check.setChecked(True)
        self.mosaic_check.stateChanged.connect(self.update_preview)
        common_aug_layout.addRow("马赛克增强:", self.mosaic_check)
        
        augmentation_layout.addWidget(common_aug_group)
        
        # HSV增强设置
        hsv_aug_group = QGroupBox("HSV增强")
        hsv_aug_layout = QGridLayout(hsv_aug_group)
        
        # 色调增强滑块
        hsv_hue_label = QLabel("色调增强:")
        self.hsv_hue_slider = QSlider(Qt.Horizontal)
        self.hsv_hue_slider.setRange(0, 100)
        self.hsv_hue_slider.setValue(15)
        self.hsv_hue_slider.valueChanged.connect(self.update_preview)
        self.hsv_hue_value = QLabel("0.015")
        
        hsv_aug_layout.addWidget(hsv_hue_label, 0, 0)
        hsv_aug_layout.addWidget(self.hsv_hue_slider, 0, 1)
        hsv_aug_layout.addWidget(self.hsv_hue_value, 0, 2)
        
        # 饱和度增强滑块
        hsv_sat_label = QLabel("饱和度增强:")
        self.hsv_sat_slider = QSlider(Qt.Horizontal)
        self.hsv_sat_slider.setRange(0, 100)
        self.hsv_sat_slider.setValue(70)
        self.hsv_sat_slider.valueChanged.connect(self.update_preview)
        self.hsv_sat_value = QLabel("0.7")
        
        hsv_aug_layout.addWidget(hsv_sat_label, 1, 0)
        hsv_aug_layout.addWidget(self.hsv_sat_slider, 1, 1)
        hsv_aug_layout.addWidget(self.hsv_sat_value, 1, 2)
        
        # 亮度增强滑块
        hsv_val_label = QLabel("亮度增强:")
        self.hsv_val_slider = QSlider(Qt.Horizontal)
        self.hsv_val_slider.setRange(0, 100)
        self.hsv_val_slider.setValue(40)
        self.hsv_val_slider.valueChanged.connect(self.update_preview)
        self.hsv_val_value = QLabel("0.4")
        
        hsv_aug_layout.addWidget(hsv_val_label, 2, 0)
        hsv_aug_layout.addWidget(self.hsv_val_slider, 2, 1)
        hsv_aug_layout.addWidget(self.hsv_val_value, 2, 2)
        
        augmentation_layout.addWidget(hsv_aug_group)
        
        # 连接滑块到标签
        self.hsv_hue_slider.valueChanged.connect(
            lambda: self.hsv_hue_value.setText(f"{self.hsv_hue_slider.value() / 1000:.3f}")
        )
        self.hsv_sat_slider.valueChanged.connect(
            lambda: self.hsv_sat_value.setText(f"{self.hsv_sat_slider.value() / 100:.1f}")
        )
        self.hsv_val_slider.valueChanged.connect(
            lambda: self.hsv_val_value.setText(f"{self.hsv_val_slider.value() / 100:.1f}")
        )
        
        tabs.addTab(augmentation_tab, "数据增强")
        
        # 高级设置标签页
        advanced_tab = QWidget()
        advanced_layout = QFormLayout(advanced_tab)
        
        # 损失函数权重
        self.box_loss_weight_spin = QDoubleSpinBox()
        self.box_loss_weight_spin.setRange(0.0, 10.0)
        self.box_loss_weight_spin.setValue(0.05)
        self.box_loss_weight_spin.setSingleStep(0.01)
        self.box_loss_weight_spin.valueChanged.connect(self.update_preview)
        advanced_layout.addRow("Box损失权重:", self.box_loss_weight_spin)
        
        self.cls_loss_weight_spin = QDoubleSpinBox()
        self.cls_loss_weight_spin.setRange(0.0, 10.0)
        self.cls_loss_weight_spin.setValue(0.5)
        self.cls_loss_weight_spin.setSingleStep(0.01)
        self.cls_loss_weight_spin.valueChanged.connect(self.update_preview)
        advanced_layout.addRow("类别损失权重:", self.cls_loss_weight_spin)
        
        self.dfl_loss_weight_spin = QDoubleSpinBox()
        self.dfl_loss_weight_spin.setRange(0.0, 10.0)
        self.dfl_loss_weight_spin.setValue(1.0)
        self.dfl_loss_weight_spin.setSingleStep(0.01)
        self.dfl_loss_weight_spin.valueChanged.connect(self.update_preview)
        advanced_layout.addRow("DFL损失权重:", self.dfl_loss_weight_spin)
        
        # 多尺度训练
        self.multi_scale_check = QCheckBox()
        self.multi_scale_check.setChecked(True)
        self.multi_scale_check.stateChanged.connect(self.update_preview)
        advanced_layout.addRow("多尺度训练:", self.multi_scale_check)
        
        # 保存频率
        self.save_period_spin = QSpinBox()
        self.save_period_spin.setRange(1, 100)
        self.save_period_spin.setValue(10)
        self.save_period_spin.valueChanged.connect(self.update_preview)
        advanced_layout.addRow("保存频率 (epochs):", self.save_period_spin)
        
        tabs.addTab(advanced_tab, "高级设置")
        
        config_layout.addWidget(tabs)
        
        # 确认按钮
        button_layout = QHBoxLayout()
        
        # 重置按钮
        reset_btn = QPushButton("重置为默认值")
        reset_btn.clicked.connect(self.reset_to_defaults)
        button_layout.addWidget(reset_btn)
        
        # 空白占位
        button_layout.addStretch()
        
        # 应用按钮
        apply_btn = QPushButton("应用设置")
        apply_btn.clicked.connect(self.apply_config)
        button_layout.addWidget(apply_btn)
        
        config_layout.addLayout(button_layout)
        
        right_panel.setWidget(config_widget)
        
        # 添加左右面板到分割器
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([300, 700])  # 设置初始大小比例
        
        main_layout.addWidget(splitter)
        
        # 初始化配置预览
        self.update_preview()
        
    def update_preview(self):
        """更新配置预览"""
        # 构建YAML格式的配置文件预览
        preview = f"""# YOLOv10 Configuration
# Generated by YOLOv10 可视化应用程序

model:
  type: {self.model_type_combo.currentText()}
  input_size: {self.input_size_combo.currentText()}

train:
  epochs: {self.epochs_spin.value()}
  batch_size: {self.batch_size_spin.value()}
  device: {self.device_combo.currentText()}
  
optimizer:
  name: {self.optimizer_combo.currentText()}
  lr: {self.lr_spin.value()}
  weight_decay: {self.weight_decay_spin.value()}
  lr_scheduler: {self.lr_scheduler_combo.currentText()}
  
augmentation:
  hsv_h: {self.hsv_hue_slider.value() / 1000:.3f}
  hsv_s: {self.hsv_sat_slider.value() / 100:.1f}
  hsv_v: {self.hsv_val_slider.value() / 100:.1f}
  flip_horizontal: {str(self.flip_h_check.isChecked()).lower()}
  flip_vertical: {str(self.flip_v_check.isChecked()).lower()}
  rotate: {str(self.rotate_check.isChecked()).lower()}
  mosaic: {str(self.mosaic_check.isChecked()).lower()}
  
loss:
  box_weight: {self.box_loss_weight_spin.value()}
  cls_weight: {self.cls_loss_weight_spin.value()}
  dfl_weight: {self.dfl_loss_weight_spin.value()}
  
advanced:
  multi_scale: {str(self.multi_scale_check.isChecked()).lower()}
  save_period: {self.save_period_spin.value()}
"""
        self.preview_text.setText(preview)
        
    def new_config(self):
        """创建新配置"""
        # 弹出对话框输入配置名称
        name, ok = QInputDialog.getText(self, "新建配置", "输入配置名称:")
        if ok and name:
            # 重置为默认设置
            self.reset_to_defaults()
            # 添加到列表
            self.config_list.addItem(name)
            # 选中新添加的项
            self.config_list.setCurrentRow(self.config_list.count() - 1)
            # 更新预览
            self.update_preview()
        
    def load_config(self):
        """加载配置文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "加载配置", "", "YAML文件 (*.yaml *.yml)"
        )
        
        if file_path:
            try:
                # 实际应用中应该解析YAML文件并更新UI控件
                import os
                config_name = os.path.basename(file_path)
                
                # 显示加载消息
                QMessageBox.information(self, "配置加载", f"已加载配置文件: {config_name}")
                
                # 添加到列表（实际应用中应检查是否已存在）
                self.config_list.addItem(config_name)
                
            except Exception as e:
                QMessageBox.critical(self, "错误", f"加载配置文件失败: {str(e)}")
        
    def save_config(self):
        """保存当前配置"""
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存配置", "", "YAML文件 (*.yaml)"
        )
        
        if file_path:
            try:
                # 获取当前配置内容
                config_content = self.preview_text.toPlainText()
                
                # 写入文件
                with open(file_path, 'w') as f:
                    f.write(config_content)
                    
                # 发出信号
                self.config_saved.emit(file_path)
                
                QMessageBox.information(self, "成功", f"配置已保存到: {file_path}")
                
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存配置文件失败: {str(e)}")
        
    def select_config(self, item):
        """选择配置列表中的项目"""
        config_name = item.text()
        
        # 在实际应用中，这里应该加载所选配置的设置
        # 简单起见，这里仅显示一个消息
        QMessageBox.information(self, "已选择配置", f"已选择配置: {config_name}\n(实际应用中将加载配置设置)")
        
    def reset_to_defaults(self):
        """重置所有设置为默认值"""
        # 基本设置
        self.model_type_combo.setCurrentText("YOLOv10s")
        self.input_size_combo.setCurrentText("640x640")
        self.batch_size_spin.setValue(16)
        self.epochs_spin.setValue(100)
        self.device_combo.setCurrentIndex(0)  # cuda:0
        
        # 优化器设置
        self.optimizer_combo.setCurrentText("SGD")
        self.lr_spin.setValue(0.01)
        self.lr_scheduler_combo.setCurrentText("cosine")
        self.weight_decay_spin.setValue(0.0005)
        
        # 数据增强设置
        self.flip_h_check.setChecked(True)
        self.flip_v_check.setChecked(False)
        self.rotate_check.setChecked(True)
        self.mosaic_check.setChecked(True)
        self.hsv_hue_slider.setValue(15)  # 0.015
        self.hsv_sat_slider.setValue(70)  # 0.7
        self.hsv_val_slider.setValue(40)  # 0.4
        
        # 高级设置
        self.box_loss_weight_spin.setValue(0.05)
        self.cls_loss_weight_spin.setValue(0.5)
        self.dfl_loss_weight_spin.setValue(1.0)
        self.multi_scale_check.setChecked(True)
        self.save_period_spin.setValue(10)
        
        # 更新预览
        self.update_preview()
        
    def apply_config(self):
        """应用当前配置"""
        # 确认是否应用
        reply = QMessageBox.question(self, "确认应用", "确定要应用当前配置吗？",
                                    QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.Yes:
            # 在实际应用中，这里应该将配置应用到系统中
            QMessageBox.information(self, "已应用", "配置已成功应用")
            # 发出信号
            self.config_saved.emit("current_config.yaml")

    def load_project_config(self, project_path):
        """加载项目的配置文件"""
        try:
            import os
            import json
            # 查找配置目录中的配置文件
            configs_dir = os.path.join(project_path, "configs")
            if not os.path.exists(configs_dir):
                return
                
            # 检查是否有配置文件
            config_files = [f for f in os.listdir(configs_dir) if f.endswith('.json')]
            if not config_files:
                return
                
            # 清空当前配置列表
            self.config_list.clear()
            
            # 加载所有配置文件
            for config_file in config_files:
                file_path = os.path.join(configs_dir, config_file)
                self.config_list.addItem(os.path.splitext(config_file)[0])
                
            # 选择第一个配置
            if self.config_list.count() > 0:
                self.config_list.setCurrentRow(0)
                self.select_config(self.config_list.item(0))
                
        except Exception as e:
            import traceback
            print(f"加载项目配置失败: {str(e)}")
            print(traceback.format_exc())


# 测试代码
if __name__ == "__main__":
    import sys
    from PyQt5.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    window = ConfigPanel()
    window.show()
    sys.exit(app.exec_()) 